{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eec680f5",
   "metadata": {},
   "source": [
    "# Agent 에 메모리(memory) 추가\n",
    "\n",
    "현재 챗봇은 과거 상호작용을 스스로 기억할 수 없어 일관된 다중 턴 대화를 진행하는 데 제한이 있습니다. \n",
    "\n",
    "이번 튜토리얼에서는 이를 해결하기 위해 **memory** 를 추가합니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "이번에는 pre-built 되어있는 `ToolNode` 와 `tools_condition` 을 활용합니다.\n",
    "\n",
    "1. [ToolNode](https://langchain-ai.github.io/langgraph/reference/prebuilt/#langgraph.prebuilt.tool_node.ToolNode): 도구 호출을 위한 노드\n",
    "2. [tools_condition](https://langchain-ai.github.io/langgraph/reference/prebuilt/#langgraph.prebuilt.tool_node.tools_condition): 도구 호출 여부에 따른 조건 분기"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b2184f4",
   "metadata": {},
   "source": [
    "\n",
    "우리의 챗봇은 이제 도구를 사용하여 사용자 질문에 답할 수 있지만, 이전 상호작용의 **context**를 기억하지 못합니다. 이는 멀티턴(multi-turn) 대화를 진행하는 능력을 제한합니다.\n",
    "\n",
    "`LangGraph`는 **persistent checkpointing** 을 통해 이 문제를 해결합니다. \n",
    "\n",
    "그래프를 컴파일할 때 `checkpointer`를 제공하고 그래프를 호출할 때 `thread_id`를 제공하면, `LangGraph`는 각 단계 후 **상태를 자동으로 저장** 합니다. 동일한 `thread_id`를 사용하여 그래프를 다시 호출하면, 그래프는 저장된 상태를 로드하여 챗봇이 이전에 중단한 지점에서 대화를 이어갈 수 있게 합니다.\n",
    "\n",
    "**checkpointing** 는 LangChain 의 메모리 기능보다 훨씬 강력합니다. (아마 이 튜토리얼을 완수하면 자연스럽게 이를 확인할 수 있습니다.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de9d9d8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b5c6228",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install -qU langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH17-LangGraph-Modules\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d2e7318",
   "metadata": {},
   "source": [
    "하지만 너무 앞서 나가기 전에, 멀티턴(multi-turn) 대화를 가능하게 하기 위해 **checkpointing**을 추가해 보도록 하겠습니다.\n",
    "\n",
    "`MemorySaver` checkpointer를 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "53d80de1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "# 메모리 저장소 생성\n",
    "memory = MemorySaver()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fc89fcf",
   "metadata": {},
   "source": [
    "**참고**\n",
    "\n",
    "이번 튜토리얼에서는 `in-memory checkpointer` 를 사용합니다. \n",
    "\n",
    "하지만, 프로덕션 단계에서는 이를 `SqliteSaver` 또는 `PostgresSaver` 로 변경하고 자체 DB에 연결할 수 있습니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51549b1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated\n",
    "from typing_extensions import TypedDict\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_teddynote.tools.tavily import TavilySearch\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.prebuilt import ToolNode, tools_condition\n",
    "\n",
    "\n",
    "########## 1. 상태 정의 ##########\n",
    "# 상태 정의\n",
    "class State(TypedDict):\n",
    "    # 메시지 목록 주석 추가\n",
    "    messages: Annotated[list, add_messages]\n",
    "\n",
    "\n",
    "########## 2. 도구 정의 및 바인딩 ##########\n",
    "# 도구 초기화\n",
    "tool = TavilySearch(max_results=3)\n",
    "tools = [tool]\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# 도구와 LLM 결합\n",
    "llm_with_tools = llm.bind_tools(tools)\n",
    "\n",
    "\n",
    "########## 3. 노드 추가 ##########\n",
    "# 챗봇 함수 정의\n",
    "def chatbot(state: State):\n",
    "    # 메시지 호출 및 반환\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "\n",
    "# 상태 그래프 생성\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "# 챗봇 노드 추가\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "\n",
    "# 도구 노드 생성 및 추가\n",
    "tool_node = ToolNode(tools=[tool])\n",
    "\n",
    "# 도구 노드 추가\n",
    "graph_builder.add_node(\"tools\", tool_node)\n",
    "\n",
    "# 조건부 엣지\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"chatbot\",\n",
    "    tools_condition,\n",
    ")\n",
    "\n",
    "########## 4. 엣지 추가 ##########\n",
    "\n",
    "# tools > chatbot\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "\n",
    "# START > chatbot\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "\n",
    "# chatbot > END\n",
    "graph_builder.add_edge(\"chatbot\", END)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c762780f",
   "metadata": {},
   "source": [
    "마지막으로, 제공된 `checkpointer`를 사용하여 그래프를 컴파일합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3d4ff857",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프 빌더 컴파일\n",
    "graph = graph_builder.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ffe3565",
   "metadata": {},
   "source": [
    "그래프의 연결성은 `LangGraph-Agent` 와 동일합니다.\n",
    "\n",
    "단지, 이번에 추가된 것은 그래프가 각 노드를 처리하면서 `State`를 체크포인트하는 것뿐입니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5622b194",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "# 그래프 시각화\n",
    "visualize_graph(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e32dbc5",
   "metadata": {},
   "source": [
    "## RunnableConfig 설정\n",
    "\n",
    "`RunnableConfig` 을 정의하고 `recursion_limit` 과 `thread_id` 를 설정합니다.\n",
    "\n",
    "- `recursion_limit`: 최대 방문할 노드 수. 그 이상은 RecursionError 발생\n",
    "- `thread_id`: 스레드 ID 설정\n",
    "\n",
    "`thread_id` 는 대화 세션을 구분하는 데 사용됩니다. 즉, 메모리의 저장은 `thread_id` 에 따라 개별적으로 이루어집니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2fea0653",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableConfig\n",
    "\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,  # 최대 10개의 노드까지 방문. 그 이상은 RecursionError 발생\n",
    "    configurable={\"thread_id\": \"1\"},  # 스레드 ID 설정\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab5c0bad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 첫 질문\n",
    "question = (\n",
    "    \"내 이름은 `테디노트` 입니다. YouTube 채널을 운영하고 있어요. 만나서 반가워요\"\n",
    ")\n",
    "\n",
    "for event in graph.stream({\"messages\": [(\"user\", question)]}, config=config):\n",
    "    for value in event.values():\n",
    "        value[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bac57c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 이어지는 질문\n",
    "question = \"내 이름이 뭐라고 했지?\"\n",
    "\n",
    "for event in graph.stream({\"messages\": [(\"user\", question)]}, config=config):\n",
    "    for value in event.values():\n",
    "        value[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "574988ab",
   "metadata": {},
   "source": [
    "이번에는 `RunnableConfig` 의 `thread_id` 를 변경한 뒤, 이전 대화 내용을 기억하고 있는지 물어보겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68ce68f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableConfig\n",
    "\n",
    "question = \"내 이름이 뭐라고 했지?\"\n",
    "\n",
    "config = RunnableConfig(\n",
    "    recursion_limit=10,  # 최대 10개의 노드까지 방문. 그 이상은 RecursionError 발생\n",
    "    configurable={\"thread_id\": \"2\"},  # 스레드 ID 설정\n",
    ")\n",
    "\n",
    "for event in graph.stream({\"messages\": [(\"user\", question)]}, config=config):\n",
    "    for value in event.values():\n",
    "        value[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "857821c4",
   "metadata": {},
   "source": [
    "## 스냅샷: 저장된 State 확인\n",
    "\n",
    "지금까지 두 개의 다른 스레드에서 몇 개의 체크포인트를 만들었습니다. \n",
    "\n",
    "`Checkpoint` 에는 현재 상태 값, 해당 구성, 그리고 처리할 `next` 노드가 포함되어 있습니다.\n",
    "\n",
    "주어진 설정에서 그래프의 `state`를 검사하려면 언제든지 `get_state(config)`를 호출하세요."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16d9c636",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableConfig\n",
    "\n",
    "config = RunnableConfig(\n",
    "    configurable={\"thread_id\": \"1\"},  # 스레드 ID 설정\n",
    ")\n",
    "# 그래프 상태 스냅샷 생성\n",
    "snapshot = graph.get_state(config)\n",
    "snapshot.values[\"messages\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ef39bd5",
   "metadata": {},
   "source": [
    "`snapshot.config` 를 출력하게 설정된 config 정보를 확인할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a203fa62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 설정된 config 정보\n",
    "snapshot.config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4917d010",
   "metadata": {},
   "source": [
    "`snapshot.value` 를 출력하게 지금까지 저장된 state 값을 확인할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "683a655b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 저장된 값(values)\n",
    "snapshot.values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6223078c",
   "metadata": {},
   "source": [
    "`snapshot.next` 를 출력하여 현재 시점에서 앞으로 찾아갈 **다음 노드를 확인** 할 수 있습니다.\n",
    "\n",
    "__END__ 에 도달하였기 때문에 다음 노드는 빈 값이 출력됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c38af74d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 다음 노드\n",
    "snapshot.next"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5576e3e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "snapshot.metadata[\"writes\"][\"chatbot\"][\"messages\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a040e572",
   "metadata": {},
   "source": [
    "복잡한 구조의 metadata 를 시각화하기 위해 `display_message_tree` 함수를 사용합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe2dbdcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.messages import display_message_tree\n",
    "\n",
    "# 메타데이터(tree 형태로 출력)\n",
    "display_message_tree(snapshot.metadata)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
